from typing import Optional

from langchain.chains.openai_functions import create_structured_output_chain
from langchain.chat_models import ChatOpenAI
from langchain.pydantic_v1 import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate


class Person(BaseModel):
    """Identifying information about a person."""
    name: str = Field(..., description="The person's name")
    age: int = Field(..., description="The person's age")
    fav_food: Optional[str] = Field(None, description="The person's favorite food")


llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, openai_api_key="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")
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a world class algorithm for extracting information in structured formats."),
        ("human", "Use the given format to extract information from the following input: {input}"),
        ("human", "Tip: Make sure to answer in the correct format"),
    ]
)

chain = create_structured_output_chain(Person, llm, prompt, verbose=True)
result = chain.run("Sally is 13 years old and likes to eat pizza")
print(result)
